import pygame
import sys
import numpy as np

# 初始化Pygame
pygame.init()

# 设置窗口大小和标题
screen_size = (640, 640)
screen = pygame.display.set_mode(screen_size)
pygame.display.set_caption("鼠标跟踪小球 - 卡尔曼滤波")

# 定义小球的属性
ball_pos = np.array([320, 320], dtype=np.float32)  # 小球初始位置，居中
ball_radius = 5  # 小球半径
ball_color = (255, 0, 0)  # 红色

# 卡尔曼滤波参数设置
dt = 0.1  # 时间步长
# 状态向量 [位置_x, 位置_y, 速度_x, 速度_y]
state = np.array([ball_pos[0], ball_pos[1], 0, 0], dtype=np.float32)

# 状态转移矩阵
A = np.array([[1, 0, dt, 0],
              [0, 1, 0, dt],
              [0, 0, 1, 0],
              [0, 0, 0, 1]], dtype=np.float32)

# 控制输入矩阵（无控制项，此处为零矩阵）
B = np.zeros((4, 4), dtype=np.float32)

# 测量矩阵，只测量位置
H = np.array([[1, 0, 0, 0],
              [0, 1, 0, 0]], dtype=np.float32)

# 预测过程噪声协方差矩阵
# 调整建议：增加 Q 的值可以提高小球对突然位置变化的灵敏度，减小则使其响应更加平稳。如果 Q 设置过大，可能导致过冲和震荡；如果 Q 设置过小，小球可能会滞后于鼠标位置变化。
Q = np.eye(4) * 0.35

# 测量噪声协方差矩阵
# 调整建议：适当增大 R 可以平滑小球的运动，但过大会使小球响应迟钝；减小 R 则使小球更贴近鼠标，但可能增加抖动。R 值设置过低时，容易导致过冲和震荡，因为小球会急于跟随测量值。
R = np.eye(2) * 5

# 初始化协方差矩阵
P = np.eye(4)

# 主循环
clock = pygame.time.Clock()
while True:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            pygame.quit()
            sys.exit()

    # 获取鼠标位置作为测量值
    mouse_x, mouse_y = pygame.mouse.get_pos()
    measurement = np.array([mouse_x, mouse_y], dtype=np.float32)

    # 卡尔曼预测步骤
    state = A @ state  # 预测状态
    P = A @ P @ A.T + Q  # 预测协方差

    # 卡尔曼更新步骤
    y = measurement - (H @ state)  # 计算测量残差
    S = H @ P @ H.T + R  # 计算残差协方差
    K = P @ H.T @ np.linalg.inv(S)  # 计算卡尔曼增益
    state = state + K @ y  # 更新状态
    P = (np.eye(4) - K @ H) @ P  # 更新协方差矩阵

    # 更新小球位置
    ball_pos[0] = state[0]
    ball_pos[1] = state[1]

    # 限制小球在屏幕内移动
    ball_pos[0] = max(ball_radius, min(screen_size[0] - ball_radius, ball_pos[0]))
    ball_pos[1] = max(ball_radius, min(screen_size[1] - ball_radius, ball_pos[1]))

    # 绘制
    screen.fill((255, 255, 255))  # 白色背景
    pygame.draw.circle(screen, ball_color, (int(ball_pos[0]), int(ball_pos[1])), ball_radius)
    pygame.display.flip()

    # 控制帧率
    clock.tick(144)

